from core.function import Function
from core.cuda import cuda
import utils

class Mul(Function):
    def forward(self, x0, x1):
        """
        执行乘法操作的前向传播。

        Parameters:
            x0 (Variable): 第一个输入变量。
            x1 (Variable): 第二个输入变量。

        Returns:
            Variable: 乘法操作的结果。
        """
        x0 = cuda.to_array(x0)
        x1 = cuda.to_array(x1)
        y = x0 * x1
        return y

    def backward(self, gy):

        """
        执行乘法操作的反向传播。

        Parameters:
            gy (Variable): 相对于前向传播结果的梯度。

        Returns:
            Tuple[Variable, Variable]: 相对于输入变量的梯度。
        """
        x0, x1 = self.inputs
        gx0 = gy * x1
        gx1 = gy * x0
        if x0.shape != x1.shape:
            gx0 = utils.functions_collect.sum_to(gx0, x0.shape)
            gx1 = utils.functions_collect.sum_to(gx1, x1.shape)
        return gx0, gx1

    def mul(self, x0, x1):
        """
        执行乘法操作。

        Parameters:
            x0 (Variable): 第一个输入变量。
            x1 (Variable): 第二个输入变量。

        Returns:
            Variable: 乘法操作的结果。
        """
        return self(x0, x1)

